/*
 * Equilibrator.cpp
 *
 *  Created on: 2 Aug 2011
 *      Author: Allan
 */

#include "Equilibrator.h"

// C++ includes
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <set>
#include <tuple>

// GeoReact includes
#include "Assemble.h"
#include "Multiphase.h"
#include "ReactionSystem.h"
#include "Utils.h"
#include "Utilities.h"

Equilibrator::Equilibrator() :
T(25.0), P(1.0)
{}

Equilibrator::Equilibrator(const vector<string>& allSpecies, const vector<string>& equilibriumSpecies, const vector<string>& equilibriumConditions) :
T(25.0), P(1.0)
{
	Initialiaze(allSpecies, equilibriumSpecies, equilibriumConditions);
}

Equilibrator::Equilibrator(const vector<string>& species) :
T(25.0), P(1.0)
{
	Initialiaze(species);
}

Equilibrator::Equilibrator(const Multiphase& multiphase, const ReactionSystem& equilibriumReactions) :
T(25.0), P(1.0)
{
	this->multiphase = multiphase;
	this->equilibriumReactions = equilibriumReactions;
	
	this->equilibriumReactions.SetIndexes(multiphase.GetSpeciesMap());
	
	n = VectorXd::Constant(multiphase.GetSpecies().size(), 1.0E-7);
}

void Equilibrator::Initialiaze(const vector<string>& allSpecies, const vector<string>& equilibriumSpecies, const vector<string>& equilibriumConditions)
{
	multiphase = AssembleMultiphase(allSpecies);
	
	equilibriumReactions = AssembleReactionSystem(equilibriumSpecies);
	
	equilibriumReactions.SetIndexes(multiphase.GetSpeciesMap());
	
	this->equilibriumConditions = equilibriumConditions;
}

void Equilibrator::Initialiaze(const vector<string>& species)
{
	Initialiaze(species, {});
}

void Equilibrator::Initialiaze(const vector<string>& species, const vector<string>& excluded)
{
	vector<string> allSpecies = AssembleSpecies(species, {"aqueous"});
	
	auto isExcluded = [&](const string& s) { return std::find(excluded.begin(), excluded.end(), s) != excluded.end(); };
	
	allSpecies.erase(std::remove_if(allSpecies.begin(), allSpecies.end(), isExcluded), allSpecies.end());
	
	const vector<string> equilibriumSpecies = RemoveSpecies(allSpecies, {"basis"});
	
	multiphase = AssembleMultiphase(allSpecies);
	
	equilibriumReactions = AssembleReactionSystem(equilibriumSpecies);
	
	equilibriumReactions.SetIndexes(multiphase.GetSpeciesMap());
	
	n = VectorXd::Constant(allSpecies.size(), 1.0E-7);
}

void Equilibrator::SetTemperaturePressure(double T, double P)
{
	this->T = T;
	this->P = P;
}

void Equilibrator::AddCondition(const string& type)
{
	eConditions.push_back( EquilibriumCondition(multiphase, equilibriumReactions) );
	
	if(type == "charge-balanced")
		eConditions.back().SetChargeBalance();
	
	else cerr << "Error: The equilibrium condition " << type << " is not valid." << endl;
}

void Equilibrator::AddCondition(const string& type, double value)
{
	eConditions.push_back( EquilibriumCondition(multiphase, equilibriumReactions) );
	
	if(type == "pH")
	{
		eConditions.back().SetAcidity(value);
		
		n[ multiphase.GetSpeciesMap()["H[+]"] ] = std::pow(10.0, -value);
	}
	
	else cerr << "Error: The equilibrium condition " << type << " is not valid." << endl;
}

void Equilibrator::AddCondition(const string& type, const string& species, double value)
{
	eConditions.push_back( EquilibriumCondition(multiphase, equilibriumReactions) );
	
	if(type == "concentration")
	{
		eConditions.back().SetComposition(species, value);
		
		n[ multiphase.GetSpeciesMap()[species] ] = value;
	}
	
	else if(type == "total-concentration")
	{
		eConditions.back().SetTotalComposition(species, value);
		
		n[ multiphase.GetSpeciesMap()[species] ] = (value > 0.0) ? value : 1.0E-7;
	}
	
	else if(type == "activity")
	{
		eConditions.back().SetActivity(species, value);
		
		n[ multiphase.GetSpeciesMap()[species] ] = value;
	}
	
	else if(type == "partial-pressure")
	{
		eConditions.back().SetPartialPressure(species, value);
		
		n[ multiphase.GetSpeciesMap()[species] ] = value/P;
	}
	
	else cerr << "Error: The equilibrium condition " << type << " is not valid." << endl;
}

void Equilibrator::ClearConditions()
{
	eConditions.clear();
	
	n = VectorXd::Constant(multiphase.GetSpecies().size(), 1.0E-7);
}

const double Equilibrator::Get(const string& quantity) const
{
	if(quantity == "ionic-strength")
	{
		auto species = multiphase.GetSpecies();
		
		double iH2O = multiphase.GetSpeciesMap()["H2O"];
		
		double I = 0.0;
		
		for(unsigned i = 0; i < species.size(); ++i)
			I += ElectricalCharge(species[i]) * ElectricalCharge(species[i]) * (55.508 * n[i]/n[iH2O]);
		
		return 0.5 * I;
	}
	
	if(quantity == "pH")
	{
		double iH = multiphase.GetSpeciesMap()["H[+]"];
		
		return -std::log10(a[iH]);
	}
	
	cerr << "Error: The quantity type " << quantity << " is not valid." << endl; return 0.0;
}

const double Equilibrator::Get(const string& quantity, const string& species) const
{
	const Index ispecies = multiphase.GetSpeciesMap()[species];
	
	if(quantity == "concentration")
		return c[ispecies];
	
	if(quantity == "activity")
		return a[ispecies];
	
	if(quantity == "activity-coeff")
		return g[ispecies];
	
	cerr << "Error: The quantity type " << quantity << " is not valid." << endl; return 0.0;
}

void Equilibrator::GetState(double&T, double& P, VectorXd& n)
{
	T = this->T;
	P = this->P;
	n = this->n;
}

void Equilibrator::Solve()
{
	n = GuessSpeciesComposition(equilibriumConditions, multiphase, T, P);
	
	Solve(T, P, n);
}

void Equilibrator::Solve(double T, double P, VectorXd& n)
{
	vector<EquilibriumCondition> eConditions = ConvertEquilibriumConditions(equilibriumConditions, multiphase, equilibriumReactions);
	
	EquilibriumSolver equilibriumSolver(multiphase, equilibriumReactions);
	
	equilibriumSolver.SetEquilibriumConditions(eConditions);
	
	equilibriumSolver.Solve(T, P, n);
	
	this->T = T;
	this->P = P;
	this->n = n;
}

void Equilibrator::Solve(double T, double P)
{
	EquilibriumSolver equilibriumSolver(multiphase, equilibriumReactions);
	
	equilibriumSolver.SetEquilibriumConditions(eConditions);
	
	equilibriumSolver.Solve(T, P, n);
	
	this->T = T;
	this->P = P;
	this->n = n;
	
	c = n;
	a = multiphase.Activities(T, P, n);
	g = multiphase.ActivityCoefficients(T, P, n);
}

void Equilibrator::PrintState() const
{
	VectorXd a = multiphase.Activities(T, P, n);
	
	cout << "Multiphase System: " << endl;
	cout << multiphase << endl << endl;
	
	cout << "Equilibrium-Controlled Reactions: " << endl;
	cout << equilibriumReactions << endl << endl;
	
	cout << "Temperature: " << T << " celsius" << endl;
	cout << "Pressure: " << P << " bar" << endl << endl;
	
	multiphase.PrintState(T, P, n);
	
	cout << endl << endl;
	
	equilibriumReactions.PrintState(T, P, n, a);
}














//void ReadThermoConditions(double& T, double& P, const string& filename)
//{
//	ifstream file(filename);
//	
//	string line;
//	
//	while(getline(file, line))
//	{
//		if(line == "ThermoConditions:")
//		{
//			getline(file, line);
//			
//			vector<string> lineT = SplitLineAndClean(line, " =");
//			
//			getline(file, line);
//			
//			vector<string> lineP = SplitLineAndClean(line, " =");
//			
//			T = atof(lineT[1].c_str());
//			P = atof(lineP[1].c_str());
//			
//			break;
//		}
//	}
//}
//
//void ReadEquilibriumConditions(vector<string>& equilibriumConditions, const string& filename)
//{
//	ifstream file(filename);
//	
//	string line;
//	
//	while(getline(file, line))
//	{
//		if(line == "EquilibriumConditions:")
//		{
//			while(getline(file, line))
//			{
//				vector<string> words = SplitLineAndClean(line, "");
//				
//				if(words.empty()) break;
//				
//				equilibriumConditions.push_back(line);
//			} 
//			
//			break;
//		}
//	}
//}
//
//void ReadAdditionalSpecies(vector<string>& additionalSpecies, const string& filename)
//{
//	ifstream file(filename);
//	
//	string line;
//	
//	while(getline(file, line))
//	{
//		if(line == "AdditionalSpecies:")
//		{
//			getline(file, line);
//			
//			additionalSpecies = SplitLineAndClean(line, ", ");
//			
//			break;
//		}
//	}
//}
//
//void CollectKnownSpecies(vector<string>& knownSpecies, const vector<string>& equilibriumConditions, const vector<string>& additionalSpecies)
//{
//	set<string> knownSpeciesSet(additionalSpecies.begin(), additionalSpecies.end());
//	
//	BOOST_FOREACH(const string& condition, equilibriumConditions)
//	{
//		vector<string> words = SplitLineAndClean(condition, " =:");
//		
//		if(words[0] == "pH" || words[0] == "ChargeBalanced")
//			continue;
//		
//		knownSpeciesSet.insert(words[1]);
//	}
//	
//	knownSpecies.assign(knownSpeciesSet.begin(), knownSpeciesSet.end());
//}

const vector<EquilibriumCondition> EquilibriumConditions(const vector<string>& conditions, const Multiphase& multiphase, const ReactionSystem& eReactions)
{
	vector<EquilibriumCondition> eConditions;
	
	BOOST_FOREACH(const string& condition, conditions)
	{
		EquilibriumCondition eCondition(multiphase, eReactions);
	
		vector<string> words = SplitLineAndClean(condition, ":= ");
	
		if(words[0] == "ChargeBalanced")
			eCondition.SetChargeBalance();
	
		else
		if(words[0] == "Composition")
		{
			const string species = words[1]; 
			const double moles = atof(words[2].c_str()); 
			
			eCondition.SetComposition(species, moles);
		}
	
		else
		if(words[0] == "TotalComposition")
		{
			const string species = words[1];
			const double moles = atof(words[2].c_str());
		
			eCondition.SetTotalComposition(species, moles);
		}
	
		else
		if(words[0] == "Activity")
		{
			const string species = words[1];
			const double activity = atof(words[2].c_str());
	
			eCondition.SetActivity(species, activity);
		}
	
		else
		if(words[0] == "pH")
		{
			const double pH = atof(words[1].c_str());
	
			eCondition.SetAcidity(pH);
		}
	
		else
		if(words[0] == "PartialPressure")
		{
			const string gas = words[1];
			const double Pg = atof(words[2].c_str());
	
			eCondition.SetPartialPressure(gas, Pg);
		}
	
		else cerr << "Error: The specified equilibrium condition " << condition << " is incorrect or invalid." << endl;
		
		eConditions.push_back(eCondition);
	}
	
	return eConditions;
}

const VectorXd InitialGuessEquilibriumComposition(const vector<string>& conditions, double P, double T, const Multiphase& multiphase)
{
	VectorXd n = VectorXd::Constant(multiphase.GetSpecies().size(), 1.0E-7);
	
	const SpeciesMap& speciesMap = multiphase.GetSpeciesMap();
	
	BOOST_FOREACH(const string& condition, conditions)
	{
		vector<string> words = SplitLineAndClean(condition, ":= ");
	
		if(words[0] == "Composition")
		{
			const string species = words[1]; 
			const double moles = atof(words[2].c_str()); 
			
			n[speciesMap[species]] = moles;
		}
	
		if(words[0] == "TotalComposition")
		{
			const string species = words[1];
			const double moles = atof(words[2].c_str());
		
			n[speciesMap[species]] = (moles > 0) ? moles : 1.0E-7;
		}
	
		if(words[0] == "Activity")
		{
			const string species = words[1];
			const double activity = atof(words[2].c_str());
	
			n[speciesMap[species]] = activity;
		}
	
		if(words[0] == "pH")
		{
			const double pH = atof(words[1].c_str());
	
			n[speciesMap["H[+]"]] = std::pow(10.0, -pH);
		}
	
		if(words[0] == "PartialPressure")
		{
			const string gas = words[1];
			const double Pg = atof(words[2].c_str());
	
			n[speciesMap[gas]] = Pg/P;
		}
	}
	
	return n;
}
